#include "viml/ast.h"
#include "viml/ast.h"
#include "viml/json_encoder.h"
#include <sstream>
#include <iostream>

namespace viml {

// NumberLiteral
void NumberLiteral::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string NumberLiteral::to_string() const {
    return std::to_string(value);
}

std::string NumberLiteral::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("NumberLiteral")},
        {"value", JSONEncoder::number_value(value)}
    });
}

// StringLiteral
void StringLiteral::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string StringLiteral::to_string() const {
    return "\"" + value + "\"";
}

std::string StringLiteral::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("StringLiteral")},
        {"value", JSONEncoder::string_value(value)}
    });
}

// Identifier
void Identifier::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string Identifier::to_string() const {
    return name;
}

std::string Identifier::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("Identifier")},
        {"name", JSONEncoder::string_value(name)}
    });
}

// SpecialConstantNode
void SpecialConstantNode::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string SpecialConstantNode::to_string() const {
    return constant_name;
}

std::string SpecialConstantNode::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("SpecialConstantNode")},
        {"constant_name", JSONEncoder::string_value(constant_name)}
    });
}

// BinaryExpression
void BinaryExpression::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string BinaryExpression::to_string() const {
    return "(" + left->to_string() + " " + operator_to_string(op) + " " + right->to_string() + ")";
}

std::string BinaryExpression::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("BinaryExpression")},
        {"operator", JSONEncoder::string_value(operator_to_string(op))},
        {"left", left->to_json()},
        {"right", right->to_json()}
    });
}

std::string BinaryExpression::operator_to_string(Operator op) {
    switch (op) {
        case Operator::ADD: return "+";
        case Operator::SUB: return "-";
        case Operator::MUL: return "*";
        case Operator::DIV: return "/";
        case Operator::MOD: return "%";
        case Operator::EQ: return "==";
        case Operator::NE: return "!=";
        case Operator::LT: return "<";
        case Operator::LE: return "<=";
        case Operator::GT: return ">";
        case Operator::GE: return ">=";
        case Operator::AND: return "&&";
        case Operator::OR: return "||";
        default: return "?";
    }
}

// UnaryExpression
void UnaryExpression::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string UnaryExpression::to_string() const {
    return "(" + operator_to_string(op) + operand->to_string() + ")";
}

std::string UnaryExpression::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("UnaryExpression")},
        {"operator", JSONEncoder::string_value(operator_to_string(op))},
        {"operand", operand->to_json()}
    });
}

std::string UnaryExpression::operator_to_string(Operator op) {
    switch (op) {
        case Operator::MINUS: return "-";
        case Operator::NOT: return "!";
        default: return "?";
    }
}

// CallExpression
void CallExpression::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string CallExpression::to_string() const {
    std::ostringstream oss;
    oss << callee->to_string() << "(";
    for (size_t i = 0; i < arguments.size(); ++i) {
        if (i > 0) oss << ", ";
        oss << arguments[i]->to_string();
    }
    oss << ")";
    return oss.str();
}

std::string CallExpression::to_json() const {
    std::vector<std::string> arg_jsons;
    for (const auto& arg : arguments) {
        arg_jsons.push_back(arg->to_json());
    }
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("CallExpression")},
        {"callee", callee->to_json()},
        {"arguments", JSONEncoder::array(arg_jsons)}
    });
}

// ListLiteral
void ListLiteral::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string ListLiteral::to_string() const {
    std::ostringstream oss;
    oss << "[";
    for (size_t i = 0; i < elements.size(); ++i) {
        if (i > 0) oss << ", ";
        oss << elements[i]->to_string();
    }
    oss << "]";
    return oss.str();
}

std::string ListLiteral::to_json() const {
    std::vector<std::string> elem_jsons;
    for (const auto& elem : elements) {
        elem_jsons.push_back(elem->to_json());
    }
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("ListLiteral")},
        {"elements", JSONEncoder::array(elem_jsons)}
    });
}

// IndexExpression
void IndexExpression::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string IndexExpression::to_string() const {
    return object->to_string() + "[" + index->to_string() + "]";
}

std::string IndexExpression::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("IndexExpression")},
        {"object", object->to_json()},
        {"index", index->to_json()}
    });
}

// DictLiteral
void DictLiteral::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string DictLiteral::to_string() const {
    std::ostringstream oss;
    oss << "{";
    for (size_t i = 0; i < pairs.size(); ++i) {
        if (i > 0) oss << ", ";
        oss << pairs[i].first->to_string() << ": " << pairs[i].second->to_string();
    }
    oss << "}";
    return oss.str();
}

std::string DictLiteral::to_json() const {
    std::vector<std::string> pair_jsons;
    for (const auto& pair : pairs) {
        pair_jsons.push_back(JSONEncoder::object({
            {"key", pair.first->to_json()},
            {"value", pair.second->to_json()}
        }));
    }
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("DictLiteral")},
        {"pairs", JSONEncoder::array(pair_jsons)}
    });
}

// LetStatement
void LetStatement::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string LetStatement::to_string() const {
    return "let " + name + " = " + value->to_string();
}

std::string LetStatement::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("LetStatement")},
        {"name", JSONEncoder::string_value(name)},
        {"value", value->to_json()}
    });
}

// EchoStatement
void EchoStatement::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string EchoStatement::to_string() const {
    return "echo " + expression->to_string();
}

std::string EchoStatement::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("EchoStatement")},
        {"expression", expression->to_json()}
    });
}

// IfStatement
void IfStatement::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string IfStatement::to_string() const {
    std::ostringstream oss;
    oss << "if " << condition->to_string() << "\n";
    for (const auto& stmt : then_body) {
        oss << "  " << stmt->to_string() << "\n";
    }
    if (!else_body.empty()) {
        oss << "else\n";
        for (const auto& stmt : else_body) {
            oss << "  " << stmt->to_string() << "\n";
        }
    }
    oss << "endif";
    return oss.str();
}

std::string IfStatement::to_json() const {
    std::vector<std::string> then_jsons;
    for (const auto& stmt : then_body) {
        then_jsons.push_back(stmt->to_json());
    }
    
    std::vector<std::string> else_jsons;
    for (const auto& stmt : else_body) {
        else_jsons.push_back(stmt->to_json());
    }
    
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("IfStatement")},
        {"condition", condition->to_json()},
        {"then_body", JSONEncoder::array(then_jsons)},
        {"else_body", JSONEncoder::array(else_jsons)}
    });
}

// WhileStatement
void WhileStatement::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string WhileStatement::to_string() const {
    std::ostringstream oss;
    oss << "while " << condition->to_string() << "\n";
    for (const auto& stmt : body) {
        oss << "  " << stmt->to_string() << "\n";
    }
    oss << "endwhile";
    return oss.str();
}

std::string WhileStatement::to_json() const {
    std::vector<std::string> body_jsons;
    for (const auto& stmt : body) {
        body_jsons.push_back(stmt->to_json());
    }
    
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("WhileStatement")},
        {"condition", condition->to_json()},
        {"body", JSONEncoder::array(body_jsons)}
    });
}

// ForStatement
void ForStatement::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string ForStatement::to_string() const {
    std::ostringstream oss;
    oss << "for " << variable << " in " << iterable->to_string() << "\n";
    for (const auto& stmt : body) {
        oss << "  " << stmt->to_string() << "\n";
    }
    oss << "endfor";
    return oss.str();
}

std::string ForStatement::to_json() const {
    std::vector<std::string> body_jsons;
    for (const auto& stmt : body) {
        body_jsons.push_back(stmt->to_json());
    }
    
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("ForStatement")},
        {"variable", JSONEncoder::string_value(variable)},
        {"iterable", iterable->to_json()},
        {"body", JSONEncoder::array(body_jsons)}
    });
}

// FunctionStatement
void FunctionStatement::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string FunctionStatement::to_string() const {
    std::ostringstream oss;
    oss << "function " << name << "(";
    for (size_t i = 0; i < parameters.size(); ++i) {
        if (i > 0) oss << ", ";
        oss << parameters[i];
    }
    oss << ")\n";
    for (const auto& stmt : body) {
        oss << "  " << stmt->to_string() << "\n";
    }
    oss << "endfunction";
    return oss.str();
}

std::string FunctionStatement::to_json() const {
    std::vector<std::string> param_jsons;
    for (const auto& param : parameters) {
        param_jsons.push_back(JSONEncoder::string_value(param));
    }
    
    std::vector<std::string> body_jsons;
    for (const auto& stmt : body) {
        body_jsons.push_back(stmt->to_json());
    }
    
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("FunctionStatement")},
        {"name", JSONEncoder::string_value(name)},
        {"parameters", JSONEncoder::array(param_jsons)},
        {"body", JSONEncoder::array(body_jsons)}
    });
}

// ReturnStatement
void ReturnStatement::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string ReturnStatement::to_string() const {
    if (value) {
        return "return " + value->to_string();
    }
    return "return";
}

std::string ReturnStatement::to_json() const {
    if (value) {
        return JSONEncoder::object({
            {"type", JSONEncoder::string_value("ReturnStatement")},
            {"value", value->to_json()}
        });
    } else {
        return JSONEncoder::object({
            {"type", JSONEncoder::string_value("ReturnStatement")},
            {"value", JSONEncoder::null_value()}
        });
    }
}

// CallStatement
void CallStatement::accept(ASTVisitor& visitor) {
    visitor.visit(*this);
}

std::string CallStatement::to_string() const {
    return "call " + call_expr->to_string();
}

std::string CallStatement::to_json() const {
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("CallStatement")},
        {"call_expr", call_expr->to_json()}
    });
}

// Program
void Program::accept(ASTVisitor& visitor) {
    return visitor.visit(*this);
}

std::string Program::to_string() const {
    std::ostringstream oss;
    for (size_t i = 0; i < statements.size(); ++i) {
        oss << statements[i]->to_string();
        if (i < statements.size() - 1) {
            oss << "\n";
        }
    }
    return oss.str();
}

std::string Program::to_json() const {
    std::vector<std::string> stmt_jsons;
    for (const auto& stmt : statements) {
        stmt_jsons.push_back(stmt->to_json());
    }
    
    return JSONEncoder::object({
        {"type", JSONEncoder::string_value("Program")},
        {"statements", JSONEncoder::array(stmt_jsons)}
    });
}

// ========== Clone实现 ========== //

// 叶子节点clone实现

std::unique_ptr<ASTNode> NumberLiteral::clone() const {
    auto cloned = std::make_unique<NumberLiteral>(value, line, column);
    return cloned;
}

std::unique_ptr<ASTNode> StringLiteral::clone() const {
    auto cloned = std::make_unique<StringLiteral>(value, line, column);
    return cloned;
}

std::unique_ptr<ASTNode> Identifier::clone() const {
    auto cloned = std::make_unique<Identifier>(name, line, column);
    return cloned;
}

std::unique_ptr<ASTNode> SpecialConstantNode::clone() const {
    auto cloned = std::make_unique<SpecialConstantNode>(constant_name, line, column);
    return cloned;
}

// 表达式节点clone实现

std::unique_ptr<ASTNode> BinaryExpression::clone() const {
    // 递归克隆左右子表达式
    auto cloned_left = left ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(left->clone().release())) : 
        nullptr;
    auto cloned_right = right ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(right->clone().release())) : 
        nullptr;
    
    auto cloned = std::make_unique<BinaryExpression>(
        std::move(cloned_left),
        op,
        std::move(cloned_right),
        line,
        column
    );
    return cloned;
}

std::unique_ptr<ASTNode> UnaryExpression::clone() const {
    // 递归克隆操作数
    auto cloned_operand = operand ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(operand->clone().release())) : 
        nullptr;
    
    auto cloned = std::make_unique<UnaryExpression>(
        op,
        std::move(cloned_operand),
        line,
        column
    );
    return cloned;
}

std::unique_ptr<ASTNode> CallExpression::clone() const {
    // 克隆callee
    auto cloned_callee = callee ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(callee->clone().release())) : 
        nullptr;
    
    // 克隆所有参数
    std::vector<std::unique_ptr<Expression>> cloned_args;
    for (const auto& arg : arguments) {
        if (arg) {
            cloned_args.push_back(
                std::unique_ptr<Expression>(dynamic_cast<Expression*>(arg->clone().release()))
            );
        }
    }
    
    auto cloned = std::make_unique<CallExpression>(
        std::move(cloned_callee),
        std::move(cloned_args),
        line,
        column
    );
    return cloned;
}

std::unique_ptr<ASTNode> ListLiteral::clone() const {
    // 克隆所有元素
    std::vector<std::unique_ptr<Expression>> cloned_elements;
    for (const auto& elem : elements) {
        if (elem) {
            cloned_elements.push_back(
                std::unique_ptr<Expression>(dynamic_cast<Expression*>(elem->clone().release()))
            );
        }
    }
    
    auto cloned = std::make_unique<ListLiteral>(std::move(cloned_elements));
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

std::unique_ptr<ASTNode> DictLiteral::clone() const {
    // 克隆所有key-value对
    std::vector<std::pair<std::unique_ptr<Expression>, std::unique_ptr<Expression>>> cloned_pairs;
    for (const auto& pair : pairs) {
        auto cloned_key = pair.first ? 
            std::unique_ptr<Expression>(dynamic_cast<Expression*>(pair.first->clone().release())) : 
            nullptr;
        auto cloned_value = pair.second ? 
            std::unique_ptr<Expression>(dynamic_cast<Expression*>(pair.second->clone().release())) : 
            nullptr;
        cloned_pairs.emplace_back(std::move(cloned_key), std::move(cloned_value));
    }
    
    auto cloned = std::make_unique<DictLiteral>(std::move(cloned_pairs));
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

std::unique_ptr<ASTNode> IndexExpression::clone() const {
    // 克隆object和index
    auto cloned_object = object ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(object->clone().release())) : 
        nullptr;
    auto cloned_index = index ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(index->clone().release())) : 
        nullptr;
    
    auto cloned = std::make_unique<IndexExpression>(
        std::move(cloned_object),
        std::move(cloned_index)
    );
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

// 语句节点clone实现

std::unique_ptr<ASTNode> LetStatement::clone() const {
    // 克隆value表达式
    auto cloned_value = value ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(value->clone().release())) : 
        nullptr;
    
    auto cloned = std::make_unique<LetStatement>(
        name,
        std::move(cloned_value),
        line,
        column
    );
    return cloned;
}

std::unique_ptr<ASTNode> EchoStatement::clone() const {
    // 克隆expression
    auto cloned_expr = expression ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(expression->clone().release())) : 
        nullptr;
    
    auto cloned = std::make_unique<EchoStatement>(
        std::move(cloned_expr),
        line,
        column
    );
    return cloned;
}

std::unique_ptr<ASTNode> IfStatement::clone() const {
    // 克隆condition
    auto cloned_cond = condition ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(condition->clone().release())) : 
        nullptr;
    
    // 克隆then_body
    std::vector<std::unique_ptr<Statement>> cloned_then;
    for (const auto& stmt : then_body) {
        if (stmt) {
            cloned_then.push_back(
                std::unique_ptr<Statement>(dynamic_cast<Statement*>(stmt->clone().release()))
            );
        }
    }
    
    // 克隆else_body
    std::vector<std::unique_ptr<Statement>> cloned_else;
    for (const auto& stmt : else_body) {
        if (stmt) {
            cloned_else.push_back(
                std::unique_ptr<Statement>(dynamic_cast<Statement*>(stmt->clone().release()))
            );
        }
    }
    
    auto cloned = std::make_unique<IfStatement>(
        std::move(cloned_cond),
        std::move(cloned_then),
        std::move(cloned_else)
    );
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

std::unique_ptr<ASTNode> WhileStatement::clone() const {
    // 克隆condition
    auto cloned_cond = condition ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(condition->clone().release())) : 
        nullptr;
    
    // 克隆body
    std::vector<std::unique_ptr<Statement>> cloned_body;
    for (const auto& stmt : body) {
        if (stmt) {
            cloned_body.push_back(
                std::unique_ptr<Statement>(dynamic_cast<Statement*>(stmt->clone().release()))
            );
        }
    }
    
    auto cloned = std::make_unique<WhileStatement>(
        std::move(cloned_cond),
        std::move(cloned_body)
    );
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

std::unique_ptr<ASTNode> ForStatement::clone() const {
    // 克隆iterable
    auto cloned_iter = iterable ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(iterable->clone().release())) : 
        nullptr;
    
    // 克隆body
    std::vector<std::unique_ptr<Statement>> cloned_body;
    for (const auto& stmt : body) {
        if (stmt) {
            cloned_body.push_back(
                std::unique_ptr<Statement>(dynamic_cast<Statement*>(stmt->clone().release()))
            );
        }
    }
    
    auto cloned = std::make_unique<ForStatement>(
        variable,
        std::move(cloned_iter),
        std::move(cloned_body)
    );
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

std::unique_ptr<ASTNode> FunctionStatement::clone() const {
    // 复制parameters（字符串向量）
    std::vector<std::string> cloned_params = parameters;
    
    // 克隆body
    std::vector<std::unique_ptr<Statement>> cloned_body;
    for (const auto& stmt : body) {
        if (stmt) {
            cloned_body.push_back(
                std::unique_ptr<Statement>(dynamic_cast<Statement*>(stmt->clone().release()))
            );
        }
    }
    
    auto cloned = std::make_unique<FunctionStatement>(
        name,
        std::move(cloned_params),
        std::move(cloned_body)
    );
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

std::unique_ptr<ASTNode> ReturnStatement::clone() const {
    // 克隆value（可能为nullptr）
    auto cloned_value = value ? 
        std::unique_ptr<Expression>(dynamic_cast<Expression*>(value->clone().release())) : 
        nullptr;
    
    auto cloned = std::make_unique<ReturnStatement>(std::move(cloned_value));
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

std::unique_ptr<ASTNode> CallStatement::clone() const {
    // 克隆call_expr
    auto cloned_call = call_expr ? 
        std::unique_ptr<CallExpression>(dynamic_cast<CallExpression*>(call_expr->clone().release())) : 
        nullptr;
    
    auto cloned = std::make_unique<CallStatement>(std::move(cloned_call));
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

std::unique_ptr<ASTNode> Program::clone() const {
    // 克隆所有语句
    std::vector<std::unique_ptr<Statement>> cloned_stmts;
    for (const auto& stmt : statements) {
        if (stmt) {
            cloned_stmts.push_back(
                std::unique_ptr<Statement>(dynamic_cast<Statement*>(stmt->clone().release()))
            );
        }
    }
    
    auto cloned = std::make_unique<Program>(std::move(cloned_stmts));
    cloned->line = line;
    cloned->column = column;
    return cloned;
}

} // namespace viml